
from libmyatomism import *
from math import *
import array 
import matplotlib.pyplot as plt


msLogger.setPriorityMin(Priority.ERROR)

unit = msUnitsManager.New("Angstrom Degree amu kcal/mol ps")

zmat = msZMat.New(unit).setId("zmat").set("H H 0 0.74191")

system = msSystem.New(unit).addEntity(zmat)

rotor = msLinearRigidRotor.New(unit).setMmtOfInertia(system).setSymmetryNumber(2)
rotor.getParameters().setParameter("Emax",200,unit)
trans = msTranslation.New(unit).setSystem(system)

gcoors = msGeneralizedCoordinates.New(unit)
q0 = msScalarVariable.New("Angstrom").set(0.7414, 0.1, 5, 0.1, 0.1).setId("q0")
gcoors.addVar(q0)

De   = unit.convert("eV",4.7457)
beta = unit.convert("m^-1",1.94196e10)
re   = unit.convert("m",0.74191e-10)



kinfct = msParser.New("Angstrom").setCoordinates(gcoors).setExpression("q0").setId("f=q0")
kinop = msKineticOperator.New(unit).set(gcoors,system)
kinop.addDynamicDof(zmat,0,kinfct)

#SiUnit=msUnitsManager.New("")
#J = 7
#A = unit.convert( msUnit.New("J.mol^-1" )  , hb*hb * Nav / SiUnit.convert( msUnit.New("amu.Angstrom^2") , kinop.kineticFunction() * 1 * 1 ) )
#epot=msParser.New("kcal/mol").setCoordinates(gcoors).setExpression(" 0.5 * conv * J*(J+1)/(q0 * q0) + De * ( 1 - exp( - a*( q0 - re ) ))^2").setId("f=Morse")
#epot.setConstant("De",De).setConstant("a",beta).setConstant("re",re).setConstant("pow",3).setConstant("conv",A).setConstant("J",J)

epot = msMorse.New().setCoordinates(gcoors)
epot.setConstant("De",De).setConstant("a",0.75*beta).setConstant("re",re).setConstant("pow",1.5)

qmotion = msQuantumMotion1d.New(unit).setEqOfMotion(kinop,epot)
qmotion.getParameters().setParameter("BasisSize",501,unit)
qmotion.getParameters().setParameter("Emax",De,unit)
print qmotion



qmotion.computeEigenValues()
qmotion.freezeKinValue( kinop.kineticFunction() )

xlist=[3.2,3.4,3.6,3.8,4,4.2,4.4,4.6,4.8,5,5.2,5.4,5.6,5.8,6,6.2,6.4,6.6,6.8,7]
ylist=[-1.045799562,-1.036075306,-1.028046231,-1.021549731,-1.016390199,-1.012359915,-1.009256479,-1.006895192,-1.005115979 ,-1.003785634,
-1.002796793,-1.002065035,-1.001525228,-1.001127856,-1.000835680,-1.000620936,-1.000463048,1.000346846,-1.000261178,1.000197871 ]

fig, axes = plt.subplots(2, 2)

axes[0, 0].plot(    epot.getXSample1d(0,500), epot.getYSample1d(0,500), "blue", 
	           [unit.convert("Bohr",x) for x in xlist], [epot.evaluate(20) - unit.convert("hartree",-1.000197871-y) for y in ylist],"gray"
		 )
axes[0, 0].axis((0,3.5,0,160))

rotor.computeEigenValues()
Jdist500=rotor.getOccupencyProbability(500)
Jdist3000=rotor.getOccupencyProbability(3000)
axes[1, 0].plot( Jdist500.getXSample1d(0,1000),   Jdist500.getYSample1d(0,1000),  "-",
                 Jdist3000.getXSample1d(0,1000),  Jdist3000.getYSample1d(0,1000), "-",)
axes[1, 0].set_xlabel("angular momentum number")
axes[1, 0].set_ylabel("Occupency probability")

# Now we compute the entropy and calorific capacity lists and plot them in a subplots.
# Again, the results are compared to those of the harmonic oscillator.
s=[]
Tlist  = [200, 300, 500, 800, 1200, 1500, 2000, 2500, 3000] 
for T in Tlist: 
    s.append(   trans.S(T)  + rotor.S(T) + qmotion.S(T)   + unit.convert("J/mol",R) )

cp=[]
for T in Tlist:
    cp.append(   trans.Cv(T) + rotor.Cv(T) + qmotion.Cv(T)   +  unit.convert("J/mol",R) )

axes[0, 1].plot( Tlist, s , "-")
axes[0, 1].set_xlabel("Temperature ["+unit.getTempStr()+"]")
axes[0, 1].set_ylabel("Total entropy ["+unit.getEnergyStr()+"]")

axes[1, 1].plot( Tlist, cp, "-")
axes[1, 1].set_xlabel("Temperature ["+unit.getTempStr()+"]")
axes[1, 1].set_ylabel("Total Cp ["+unit.getEnergyStr()+"]")

# this is the experimental values, coming from the janaf table,  
# see http://kinetics.nist.gov/janaf/html/H-050.html, defined in J/mol
TListexp  =  [200,     300,     500,     800,     1200,   1500,    2000,    2500,    3000   ] 
ExpListS   = [119.412, 130.858, 145.737, 159.548, 171.79, 178.846, 188.418, 196.243, 202.891]   
ExpListCp =  [27.447,  28.849,  29.260,  29.624,  30.992, 32.298,  34.280,  35.842,  37.087 ] 
 
axes[0, 1].plot( TListexp, [ unit.convert("J/mol",x) for x in ExpListS],   "r*" )
axes[1, 1].plot( TListexp, [ unit.convert("J/mol",x) for x in ExpListCp],  "r*" )

fig.set_size_inches(20,20)
fig.savefig('H2_0.png')
fig.show()

